from typing import List, Union, Optional
from Embedding import get_embp_embedding, get_embq_embedding, parser_Message
import json
from utils.logger import logger
import time
from config import Config
from openai import OpenAI

class EmbeddingService:
    def __init__(self):
        # 讯飞API配置
        self.appid = Config.XUNFEI_APPID
        self.api_secret = Config.XUNFEI_API_SECRET
        self.api_key = Config.XUNFEI_API_KEY
        self.vector_dim = Config.VECTOR_DIM
        
        # 豆包API配置
        if Config.EMBEDDING_TYPE == "douban":
            if not Config.ARK_API_KEY:
                raise ValueError("豆包 API key 未设置")
            try:
                self.douban_client = OpenAI(
                    api_key=Config.ARK_API_KEY,
                    base_url=Config.ARK_BASE_URL,
                    timeout=Config.VECTOR_REQUEST_TIMEOUT
                )
            except Exception as e:
                logger.error(f"初始化豆包客户端失败: {str(e)}")
                raise
    
    @property
    def dimension(self) -> int:
        """返回向量维度"""
        return self.vector_dim
    
    def _get_douban_embeddings(self, texts: List[str], type: str = "add_documents") -> Optional[List[List[float]]]:
        """获取豆包向量
        
        Args:
            texts: 文本列表
            type: 调用类型，可选值为 "search" 或 "add_documents"
            
        Returns:
            Optional[List[List[float]]]: 向量列表
        """
        try:
            # 限制每批次最大数量为256
            BATCH_SIZE = 128
            all_embeddings = []
            
            # 分批处理
            for i in range(0, len(texts), BATCH_SIZE):
                batch_texts = texts[i:i + BATCH_SIZE]

                # 调用豆包API
                response = self.douban_client.embeddings.create(
                    model=Config.ARK_MODEL,
                    input=batch_texts
                )
                
                # 提取向量
                batch_embeddings = [item.embedding for item in response.data]
                all_embeddings.extend(batch_embeddings)
                
                # 添加日志
                logger.info(f"已处理 {len(all_embeddings)}/{len(texts)} 个文本的向量")
                
                # 可选：添加延时避免频繁请求
                if type == "add_documents":
                    time.sleep(3)
            
            return all_embeddings
            
        except Exception as e:
            logger.error(f"获取豆包向量失败: {str(e)}")
            return None

    def get_embeddings(self, texts: Union[str, List[str]], 
                      batch_size: int = 32) -> List[List[float]]:
        """生成文本嵌入向量"""
        try:
            if isinstance(texts, str):
                texts = [texts]
            
            # 使用豆包API
            if Config.EMBEDDING_TYPE == "douban":
                return self._get_douban_embeddings(texts, "add_documents")
            
            # 使用讯飞API
            all_embeddings = []
            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i + batch_size]
                batch_embeddings = self._process_batch(batch_texts)
                all_embeddings.extend(batch_embeddings)
            
            return all_embeddings
        except Exception as e:
            logger.error(f"生成嵌入向量时发生错误: {str(e)}")
            return None
    
    def _process_batch(self, texts: List[str]) -> List[List[float]]:
        """处理一批文本"""
        start_batch = time.time()
        embeddings = []
        
        for i, text in enumerate(texts, 1):
            start = time.time()
            messages = {
                "messages": [
                    {"content": text, "role": "user"}
                ]
            }
            
            response = get_embp_embedding(
                messages,
                appid=self.appid,
                apikey=self.api_key,
                apisecret=self.api_secret
            )
            
            vector = parser_Message(response)
            if vector is not None:
                embeddings.append(vector.tolist())
            
            logger.info(f"处理第 {i}/{len(texts)} 个文本，耗时: {(time.time()-start)*1000:.2f}ms")
            time.sleep(0.5)  # API限流
        
        logger.info(f"批处理完成，共 {len(texts)} 个文本，总耗时: {(time.time()-start_batch)*1000:.2f}ms")
        return embeddings

    def get_query_embedding(self, query: str, type: str = "add_documents") -> List[float]:
        """专门用于查询文本的向量化
        
        Args:
            query: 查询文本
            type: 调用类型，可选值为 "search" 或 "add_documents"
            
        Returns:
            List[float]: 向量表示
        """
        try:
            # 使用豆包API
            if Config.EMBEDDING_TYPE == "douban":
                embeddings = self._get_douban_embeddings([query], type)
                return embeddings[0] if embeddings else None
            
            # 使用讯飞API
            messages = {
                "messages": [
                    {"content": query, "role": "user"}
                ]
            }
            
            start = time.time()
            response = get_embq_embedding(
                messages,
                appid=self.appid,
                apikey=self.api_key,
                apisecret=self.api_secret
            )
            
            vector = parser_Message(response)
            end = time.time()
            logger.info(f"查询向量化耗时: {(end-start)*1000:.2f}ms")
            
            return vector.tolist() if vector is not None else None
            
        except Exception as e:
            logger.error(f"查询向量化失败: {str(e)}")
            return None
